package com.llmops.demo;

import com.llmops.core.Context;
import com.llmops.core.Edge;
import com.llmops.core.WorkflowGraph;
import com.llmops.core.node.EndNode;
import com.llmops.core.node.RagNode;
import com.llmops.core.node.StartNode;

import java.util.Collections;
import java.util.List;
import java.util.Map;

public class Rag_Graph {

    public static void main(String[] args) {
        WorkflowGraph workflow = new WorkflowGraph("1");
        Context context = new Context();
        context.setVariable("sessionId",1);

        workflow.addNode(new StartNode("data_input","data_input",Collections.singletonMap("input","提示词")));
        workflow.addNode(new RagNode("rag_Node","rag_Node", Map.of("classifications",List.of("售前","售后"))));
        workflow.addNode(new EndNode("result_exporterA","result_exporterA", Collections.singletonMap("v", "A")));

        workflow.addEdge(new Edge("data_input", "rag_Node"));
        workflow.addEdge(new Edge("rag_Node", "result_exporterA"));


        workflow.validate();
//        workflow.execute();

        workflow.execute(context);

        System.out.println("Workflow execution completed");
    }
}
